import json
from collections import defaultdict
import networkx as nx
from typing import Dict, List, Tuple, Set

import numpy as np

import topo1
from DHTOA1 import matrix_to_communication_pairs


def allocate_switches_with_es_optimized(
        G: nx.Graph,
        T: List[Tuple[str, str]],
        D: Dict[Tuple[str, str], float],
        tor_switches: Set[str],
        optical_switches: Set[str],
        k_paths: int = 30
) -> Tuple[Dict, Dict, Dict]:
    """
    优化路径分配算法，考虑链路速率（Gbps）和流量需求（GB）。

    主要修改点：
    1. 链路容量表示传输速率（Gbps），需求表示数据大小（GB）。
    2. 延迟计算公式调整为：时间 = 数据大小 / 速率。
    """

    path_allocations = {}
    link_utilization = defaultdict(float)  # 记录链路当前已分配流量（Gbps）
    delay_info = {}

    # 记录链路原始速率（Gbps）
    original_capacity = {tuple(sorted(edge)): G.edges[edge]['capacity']
                         for edge in G.edges()}
    for edge in original_capacity:
        print(f"{edge}原始链路容量(速率)：", original_capacity[edge])

    sorted_pairs = sorted(T, key=lambda x: D[x], reverse=True)
    print("通信对排序：", sorted_pairs)

    for pair in sorted_pairs:
        s, d = pair
        if s == d:
            continue

        demand = D[pair]  # 需求数据量（GB）
        delay_info[pair] = None
        # print(f"\n=== 处理通信对 {s}→{d} [需求: {demand} GB] ===")

        try:
            paths = nx.shortest_simple_paths(G, s, d, weight='hops')
            candidate_paths = [p for _, p in zip(range(k_paths), paths)]
            # print(f"  候选路径：")
            for p in candidate_paths:
                pass
                # print(f"    {p}")
        except nx.NetworkXNoPath:
            # print(f"  路径发现失败")
            continue

        # feasible_paths = []
        # for path in candidate_paths:
        #     # valid = True
        #
        #     # if not any(n in optical_switches for n in path):
        #     #     print(f"  路径 {path} 无效：未经过光交换机")
        #     #     continue
        #
        #     for i in range(len(path) - 1):
        #         u, v = path[i], path[i + 1]
        #         edge = tuple(sorted((u, v)))
        #
        #     #     remaining_rate = original_capacity[edge] - link_utilization[edge]
        #     #     if remaining_rate <= 0:
        #     #         print(f"  链路 {edge} 速率不足")
        #     #         valid = False
        #     #         break
        #     #
        #     # if valid:
        #     #     feasible_paths.append(path)
        #     #     print(f"  可行路径: {'→'.join(path)}")
        #
        # # if not feasible_paths:
        # #     print(f"  无可行路径")
        # #     continue
        #
        # # min_delay = float('inf')
        # # best_path = None
        min_delay = float('inf')
        for path in candidate_paths:

            total_delay = 0.0
            for i in range(len(path) - 1):
                u, v = path[i], path[i + 1]
                edge = tuple(sorted((u, v)))
                # remaining_rate = original_capacity[edge] - link_utilization[edge]
                # print(f"    链路 {edge} 速率: {remaining_rate:.2f} Gbps")

                # if remaining_rate <= 0:
                #     total_delay = float('inf')
                #     break

                delay_segment = (demand+link_utilization[edge])*8 / original_capacity[edge]  # 时间 = 数据量 / 速率
                # print(f"      {edge}延迟段: {delay_segment:.2f} s")
                total_delay += delay_segment
                # print(f"      当前总延迟: {total_delay:.2f} s")

            # print(f"  路径 {'→'.join(path)} 预估延迟: {total_delay:.2f} s")

            if total_delay < min_delay:
                min_delay = total_delay
                best_path = path
            # print(f"  最佳路径 {'→'.join(best_path)} 预估延迟: {min_delay:.2f} s")

        if best_path:
            path_allocations[pair] = best_path
            for i in range(len(best_path) - 1):
                u, v = best_path[i], best_path[i + 1]
                edge = tuple(sorted((u, v)))
                # link_utilization[edge] += demand / min_delay  # 按时间分配速率
                # 更新链路上已经分配的流量
                link_utilization[edge] += demand
                delay_info[pair] = min_delay
                # print(f"  更新链路 {edge}: 当前链路分配的流量为={link_utilization[edge]:.2f} Gbps")

    return path_allocations, link_utilization,delay_info


# 修改后的构建子图方法
def find_electric_paths(G, s, d, k_paths):
    # Step 1: 创建子图并添加所有节点
    G_electric = nx.Graph()
    G_electric.add_nodes_from(G.nodes())  # 确保所有节点存在
    # print(f"构建子图electric_paths: {G_electric.nodes()}")

    # 输出拓扑中所有的边
    for u, v, attr in G.edges(data=True):
        if attr.get('type') == 'electrical':
            G_electric.add_edge(u, v, **attr)
        # print(f"边: {u}-{v}, 属性: {attr}")

    # 输出G_electric的节点和边
    # print("electric_paths:")
    # for node in G_electric.nodes():
    #     print(f"节点: {node}")
    # for u, v, attr in G_electric.edges(data=True):
    #     print(f"边: {u}-{v}, 属性: {attr}")
    # Step 3: 检查节点是否存在
    if not G_electric.has_node(s) or not G_electric.has_node(d):
        print(f"节点 {s} 或 {d} 在电链路子图中不存在")
        return []

    # Step 4: 搜索路径
    try:
        paths = nx.shortest_simple_paths(G_electric, s, d, weight='hops')
        candidate_paths = [p for _, p in zip(range(k_paths), paths)]
        return candidate_paths
    except nx.NetworkXNoPath:
        return []

def allocate_electric_with_es_optimized (
        G: nx.Graph,
        T: List[Tuple[str, str]],
        D: Dict[Tuple[str, str], float],
        init_util: Dict[Tuple[str, str], float],
        tor_switches: Set[str],
        optical_switches: Set[str],
        k_paths: int = 100
) -> Tuple[Dict, Dict, Dict]:
    """
    优化路径分配算法，考虑链路速率（Gbps）和流量需求（GB）。

    主要修改点：
    1. 链路容量表示传输速率（Gbps），需求表示数据大小（GB）。
    2. 延迟计算公式调整为：时间 = 数据大小 / 速率。
    """

    # print("===========测试============")
    # print("输出流量矩阵信息")
    print(type(D))
    # for pair,demandin in D.items():
    #     print(f"  {pair}的流量为{demandin}")
    # print("===========测试============")

    path_allocations = {}
    # link_utilization = defaultdict(float)  # 记录链路当前已分配流量（Gbps）
    link_utilization = init_util
    # 遍历通信对，如果通信对在init_util中，则初始化已分配流量为0
    # print("输出初始化链路已分配的流量信息：")
    for edge in link_utilization:
        pass
        # print(f"  初始化链路 {edge} 已使用容量为{link_utilization[edge]}")
    delay_info = {}

    # 记录链路原始速率（Gbps）
    original_capacity = {tuple(sorted(edge)): G.edges[edge]['capacity']
                         for edge in G.edges()}
    print("原始链路容量：", original_capacity)

    # sorted_pairs = sorted(T, key=lambda x: D[x], reverse=True)
    # print("通信对排序：", sorted_pairs)
    # 创建有效通信对的需求字典
    # print("输出通信对")
    for pair in T:
        pass
        # print(f"  通信对 {pair} ")
    print("输出流量矩阵的内容")
    for pair in D:
        pass
        # print(f"  流量矩阵 {pair} 的值为 {D[pair]} GB")

    valid_demand_dict = {pair: D[pair] for pair in T if pair in D}
    # print("有效通信对：", valid_demand_dict)
    for pair in valid_demand_dict:
        pass
        # print(f"有效通信对 {pair} 的需求为 {D[pair]} GB")
    # 按流量需求降序排序（仅处理有效通信对）
    sorted_items = sorted(
        valid_demand_dict.items(),
        key=lambda x: x[1],
        reverse=True
    )
    sorted_pairs = [pair for pair, _ in sorted_items]
    print("============测试================")
    print(type(sorted_pairs))
    print(sorted_pairs)
    print("============测试================")
    for pair in sorted_pairs:
        s, d = pair
        demand = D[pair]  # 需求数据量（GB）
        delay_info[pair] = None
        pass
        print(f"\n=== 处理通信对 {s}→{d} [需求: {demand} GB] ===")

        # try:
        #     paths = nx.shortest_simple_paths(G, s, d, weight='hops')
        #     candidate_paths = [p for _, p in zip(range(k_paths), paths)]
        #     print(f"  候选路径：")
        #     for p in candidate_paths:
        #         pass
        #
        #         # print(f"    {p}")
        #         # print("测试测试测试")
        # except nx.NetworkXNoPath:
        #     print(f"  路径发现失败")
        #     continue

        candidate_paths = find_electric_paths(G, s, d, k_paths)
        print("=======================测试=========================")
        for path in candidate_paths:
            print(f"  {pair}的候选路径为: {'→'.join(path)}")


        feasible_paths = []
        for path in candidate_paths:
            valid = True

            # if not any(n in optical_switches for n in path):
            #     print(f"  路径 {path} 无效：未经过光交换机")
            #     continue

            for i in range(len(path) - 1):
                u, v = path[i], path[i + 1]

                # 如果edge不是电链路，则跳过该路径
                if G[u][v]['type'] != 'electrical':
                    # print(f"  路径 {path} 无效：未经过电交换机")
                    valid = False
                    continue
                else:
                    edge = tuple(sorted((u, v)))
                #     remaining_rate = original_capacity[edge] - link_utilization[edge]
                #     if remaining_rate <= 0:
                #         print(f"  链路 {edge} 速率不足")
                #         valid = False
                #         break

            if valid:
                feasible_paths.append(path)
                # print("测试测试测试测试测试")
                print(f"  可行路径: {'→'.join(path)}")

        if not feasible_paths:
            print(f"  无可行路径")
            continue

        min_delay = float('inf')
        best_path = None

        for path in feasible_paths:
            total_delay = 0.0
            for i in range(len(path) - 1):
                u, v = path[i], path[i + 1]
                edge = tuple(sorted((u, v)))
                # remaining_rate = original_capacity[edge] - link_utilization[edge]
                # 更新链路上已经分配的流量
                # exist_traffic = link_utilization[edge]
                if edge in link_utilization:
                    exist_traffic = link_utilization[edge]
                else:
                    exist_traffic = 0.0

                # print(f"    链路 {edge} 已经分配流量: {exist_traffic:.2f} Gbps")

                # if remaining_rate <= 0:
                #     total_delay = float('inf')
                #     break

                delay_segment = (demand+exist_traffic)*8 / original_capacity[edge]  # 时间 = 数据量 / 速率
                # print(f"      延迟段:  （{demand}+{exist_traffic}）*8 / {original_capacity[edge]}={delay_segment:.2f} s")
                total_delay += delay_segment
                # print(f"      更新当前总延迟: {total_delay:.2f} s")

            # print(f"  路径 {'→'.join(path)} 预估延迟: {total_delay:.2f} s")

            if total_delay < min_delay:
                min_delay = total_delay
                best_path = path
            else:
                print(f"{total_delay} < {min_delay}")

            print(f"  最佳路径 {'→'.join(best_path)} 预估延迟: {min_delay:.2f} s")

        if best_path:
            delay_info[pair] = min_delay
            print("有best_path")
            path_allocations[pair] = best_path
            # print("==========测试测试测试===========")
            # print(f"最好的路径为{best_path}")
            for i in range(len(best_path) - 1):
                u, v = best_path[i], best_path[i + 1]
                edge = tuple(sorted((u, v)))
                if edge not in link_utilization:
                    link_utilization[edge] = demand
                else:
                    link_utilization[edge] += demand  # 按时间分配速率
                    # print(f"min_delay，{min_delay}")
                    # delay_info[pair] = min_delay
                    # print(f"Min_delay{delay_info}")

                # print(f"  更新链路 {edge}: 当前已经分配的流量为={link_utilization[edge]:.2f} Gbps")
        else:
            print("没有best_path")

    return path_allocations, link_utilization,delay_info

# 测试代码
if __name__ == "__main__":
    # G = nx.Graph()
    # edges = [
    #     ('ToR1', 'os1', 100), ('ToR2', 'os1', 100),
    #     ('ToR3', 'os2', 100), ('ToR4', 'os2', 100),
    #     ('os1', 'Core', 400), ('os2', 'Core', 400)
    # ]
    # for u, v, c in edges:
    #     G.add_edge(u, v, capacity=c)  # 速率 Gbps
    #
    # tor_switches = {'ToR1', 'ToR2', 'ToR3', 'ToR4'}
    # optical_switches = {'os1', 'os2'}
    #
    # T = [('ToR1', 'ToR3'), ('ToR2', 'ToR4'), ('ToR1', 'ToR4')]
    # D = {
    #     ('ToR1', 'ToR3'): 80,  # 80 GB
    #     ('ToR2', 'ToR4'): 60,  # 60 GB
    #     ('ToR1', 'ToR4'): 120  # 120 GB
    # }
    #
    # paths, util = allocate_switches_with_es_optimized(
    #     G, T, D, tor_switches, optical_switches, k_paths=3
    # )

    # print("\n=== 最终结果 ===")
    # print("路径分配方案:", paths)
    # print("链路利用率:", util)
    G=topo1.build_hybrid_topology_shared_tor(num_tor=4)
    # T=[('tor1', 'tor2'), ('tor1', 'tor3'), ('tor1', 'tor4'), ('tor2', 'tor3'), ('tor2', 'tor4'), ('tor3', 'tor4')]
    with open("traffic_matrices.json", 'r') as f:
        traffic_matrices = json.load(f)
    all_to_all_burst_4 = traffic_matrices['4_nodes']['all_to_all_burst']
    all_to_all_burst_4_matrix = np.array(all_to_all_burst_4)
    all_to_all_burst_4_dict = matrix_to_communication_pairs(all_to_all_burst_4_matrix, node_prefix="ToR")
    T_pair=matrix_to_communication_pairs(all_to_all_burst_4_matrix, node_prefix="ToR")
    T=[]
    for pair in T_pair:
        T.append(pair)

    tor_switches = {'ToR1', 'ToR2', 'ToR3', 'ToR4'}
    optical_switches = {'os1', 'os2'}

    paths, util ,delays= allocate_switches_with_es_optimized(
        G, T, all_to_all_burst_4_dict, tor_switches, optical_switches, k_paths=30
    )

    print("\n=== 最终结果 ===")
    for pair, path in paths.items():
        print(f"通信对 {pair}: 路径 {'→'.join(path)}")

    for pair, rate in util.items():
        print(f"链路 {pair}: 当前利用率={rate:.2f} Gbps")

    for pair, delay in delays.items():
        pass
        # status = f"{delay:.2f}s" if delay not in [None, float('inf')] else "超时/未分配"
        # print(f"通信对 {pair}: {status}")


